# ShaBa - Set of utilies for easy backups.
# Copyright (C) 2012 Damien Vincent
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

import argparse
import hashlib
import os.path
import os
import signal
import sys

FILENAME_SHA1_SIZE = 40

stop_backup = 0
def sigint_handler(signal, frame):
  sys.stderr.write('Interrupting integrity check...\n')
  stop_backup = 1

# Calculate the hash of the content of a file
def compute_hash_file(f, hasher, blocksize=65536):
  buf = f.read(blocksize)
  while len(buf) > 0:
    hasher.update(buf)
    buf = f.read(blocksize)
  return hasher.digest()

# Integrity of one file
def integrity_one_file(filename):
  try:
    f = file(filename, 'r')
  except IOError:
    return (0, [])
  f_hash = compute_hash_file(f, hashlib.sha1())
  hashname = ''.join( [ "%02x" % ord( x ) for x in f_hash ] )
  filename_base = os.path.basename(filename)
  cmp_result = (hashname == filename_base)
  return (cmp_result, f_hash)

# Verify that the directory is a directory that contains only hash files.
# => No subdirectories or other files.
def is_hash_dir(dirname):
  for root, sub_folders, files in os.walk(dirname):
    if sub_folders:
      return 0
    for f in files:
      if len(f) != FILENAME_SHA1_SIZE:
        return 0
  return 1

# Recursively calculate the total size of files inside |dirname|.
def backup_total_size(dirname):
  byte_count = 0
  for root, sub_folders, files in os.walk(dirname):
    for f in files:
      filename_source = os.path.join(root, f)
      if not os.path.islink(filename_source):
        byte_count += os.path.getsize(filename_source)
  return byte_count

# Verify the integrity of a hash directory
def integrity(backup_dir_hash, backup_size):
  byte_count = 0

  for root, sub_folders, files in os.walk(backup_dir_hash):
    root_relative = os.path.relpath(root, backup_dir_hash)
    for f in files:
      if stop_backup:
        return
      filename_relative = os.path.join(root_relative, f)
      progress = (100 * byte_count) / backup_size
      filename_display = filename_relative + (" " * 80)
      sys.stderr.write((('\r[%3d%%]: ' % progress) + filename_display[0:79]))
      filename_hash = os.path.join(backup_dir_hash, root_relative, f)
      if os.path.islink(filename_hash):
        pass
      else:
        (r, f_hash) = integrity_one_file(filename_hash)
        if r == 0:
          sys.stderr.write('\n')
          sys.stderr.write(('file: ' + filename_hash + ' is corrupted\n'))
          hashname = ''.join( [ "%02x" % ord( x ) for x in f_hash ] )
          sys.stderr.write(('Its hash is now: ' + hashname + '\n'))
        byte_count += os.path.getsize(filename_hash)
  sys.stderr.write('\n')

def main():
  signal.signal(signal.SIGINT, sigint_handler)

  parser = argparse.ArgumentParser(description='Verify the integrity of a backup.')
  parser.add_argument('--hash', action="store", dest="backup_dir_hash",
                      help='Hash directory')
  args = parser.parse_args()

  backup_dir_hash   = args.backup_dir_hash

  sys.stderr.write('Step #1: Checking directories\n')
  if not is_hash_dir(backup_dir_hash):
    sys.stderr.write('Error: Hash directory contains some sub directories or common files\n')
    sys.exit(1)

  sys.stderr.write('Step #2: Collecting backup statistics\n')
  backup_size = backup_total_size(backup_dir_hash)
  sys.stderr.write(('Total size to verify: %d bytes\n' % backup_size))

  sys.stderr. write('Step #3: Verifying integrity\n')
  integrity(backup_dir_hash, backup_size)

if __name__ == "__main__":
  main()

